/*
 * @Author: LibraStalker 1299640057@qq.com
 * @Date: 2023-01-27 23:44:51
 * @FilePath: facedetector.ts
 * @Description: 人脸框检测
 */
import {Runner} from '@paddlejs/paddlejs-core';
import {ModelSDK} from "./base";
import {FaceInfo, FeedShape, Location, ModelConfig, NaturalSize, all_supported_face_field} from "../commons/interface";

/**
 * 人脸框检测
 */
export class FaceDetector extends ModelSDK {
    readonly defaultFeedShape: FeedShape = {
        fc: 3,
        fw: 1024,
        fh: 1024
    };

    readonly defaultModelConfig: ModelConfig = {
        modelPath: 'https://paddlejs.cdn.bcebos.com/models/fuse/facedetect_opt/model.json',
        mean: [0.407843137, 0.458823529, 0.482352941],
        std: [0.5, 0.5, 0.5],
        feedShape: this.defaultFeedShape
    };

    runner: Runner = null;
    inputSize = {} as NaturalSize;
    scale: number;
    // 需要使用的额外检测负载
    face_field: [string, ModelSDK][]; // ModelSDK的意思是ModelSDK的子类的实例，暂时还不知道怎么用泛型来表示类型所以这样表示了。

    /**
     * 初始化 @member modelConfig 和 @member face_field
     * @param modelConfig 模型配置
     * @param face_field 需要使用的额外负载
     */
    constructor(modelConfig?: ModelConfig, face_field: [string, ModelConfig?][] = []) {
        super(modelConfig);
        super.__init__();
        face_field.forEach(([field_string, _]) => {
            if (Object.keys(all_supported_face_field).indexOf(field_string) < 0) {
                throw Error('不支持的face_field: ' + field_string);
            }
        });
        this.face_field = face_field.map(([field_string, model_config]) => {
            return [field_string, new all_supported_face_field[field_string](model_config)];
        });
    }

    /**
     * 初始化当前模型和所有的额外负载，在当前模型和所有额外负载都初始化完成后才会resolve
     * @returns Promise<any>
     */
    async init(): Promise<any> {
        const init_promises = [super.init()];
        this.face_field.map(async ([_, model_SDK_instance]) => {
            init_promises.push(model_SDK_instance.init());
        });
        // 参考packages/paddlejs-models/gesture/src/index.ts里的写法，这里必须把所有的初始化用的Promise放在一起用Promise.all()，
        // 如果await一个模型初始化完成后再初始化另一个模型就会在浏览器控制台出现类似的警告：
        // 27WebGL: INVALID_OPERATION: framebufferTexture2D: object does not belong to this context
        // 27WebGL: INVALID_OPERATION: useProgram: object does not belong to this context
        // 39WebGL: INVALID_OPERATION: bindTexture: object does not belong to this context
        // 39WebGL: INVALID_OPERATION: uniform1i: location not for current program
        // 参考：https://stackoverflow.com/questions/29059306/three-js-webgl-invalid-operation-bindtexture-object-not-from-this-context/29067763
        // 原因可能是初始化多个模型导致gl的texture或context出现冲突
        return await Promise.all(init_promises);
        // return new Promise(() => this);
    }

    /**
     * 从图片中检测所有的人脸，并对所有检测出的人脸区域使用检测负载进行再次检测，并添加到结果中
     * @param input 原始图片，可以通过createImageBitmap()获取
     * @param opts 可选参数
     * @param opts.threshold 取值范围是0-1，默认0.6，表示只取置信度大于此数值的人脸
     * @param opts.shrink 取值范围是0-1，默认0.4，预测前先将图片按此比例缩小，对于占图片比例较大的人脸可以提升检测效果
     * @returns 每个人脸识别的结果，包含了所有请求负载，返回的格式参考https://ai.baidu.com/sdk
     */
    async detect(
        input: ImageBitmap,
        opts?
    ): Promise<FaceInfo[]> {
        this.inputSize.naturalWidth = input.width;
        this.inputSize.naturalHeight = input.height;
        const {threshold = 0.6, shrink = 0.4} = Object.assign({}, opts);

        const inputFeed = this.preprocess(input, shrink); // 预处理
        const dataOut = await this.runner.predictWithFeed(inputFeed); // 预测
        // this.scale = Math.min(this.inputSize.naturalWidth, this.inputSize.naturalHeight) * shrink;
        // const dataOut = [[1, 0.999731719493866, 0.4461297392845154, 0.46200433373451233, 0.5182010531425476, 0.5515674352645874]];
        // [1, 0.9133170247077942, 0.4231080412864685, 0.43493539094924927, 0.5277166962623596, 0.563701868057251]
        //  [1, 0.9997240304946899, 0.44611501693725586, 0.46211671829223633, 0.5182320475578308, 0.5516325235366821]
        //  [1, 0.999731719493866, 0.4461297392845154, 0.46200433373451233, 0.5182010531425476, 0.5515674352645874]
        console.debug(dataOut, threshold); // 左上右下
        const face_list = this.postprocess(dataOut, threshold); // 后处理
        // 如果对face_list用map会导致人脸关键点检测并发执行，而this.runner并不是异步安全的，所以会导致所有人脸识别出的结果相同，所以改成了for循环
        for (const i in face_list) {
            if (!isNaN(Number(i))) {
                const face_info = face_list[i];
                // 获取原始图片中人脸检测框框出的部分，即仅包含人脸的图片
                const image_bitmap = await createImageBitmap(input, face_info.location.left, face_info.location.top,
                    face_info.location.width, face_info.location.height);
                // 单个人脸的所有检测负载组合而成的Promise
                const one_face_tasks = this.face_field.map(async ([field_string, model_SDK_instance]) => {
                    console.debug(face_info, image_bitmap);
                    face_info[field_string] = await model_SDK_instance.detect(image_bitmap, face_info.location);
                });
                await Promise.all(one_face_tasks); //等待单个人脸检测完成
            }
        }
        // await Promise.all(all_tasks);
        console.log(face_list);
        return face_list;
    }

    /**
     * 预处理，将图片按shrink比例缩小，然后按比例缩放到feedShape，若长宽比与feedShape不同则多余的部分留空，
     * @param input 原始图片，可以通过createImageBitmap()获取
     * @param shrink 取值范围是0-1，图片将按此比例缩小，对于占图片比例较大的人脸可以提升检测效果
     * @returns InputFeed[] 可以直接传给Runner.predictWithFeed
     */
    preprocess(input: ImageBitmap, shrink: number) {
        // shrink --> scale
        const scale = this.scale = Math.min(this.inputSize.naturalWidth, this.inputSize.naturalHeight) * shrink;
        return this.runner.mediaProcessor.process(
            input,
            Object.assign({}, this.modelConfig, {scale}),
            this.feedShape
        );
    }

    /**
     * 后处理，将模型输出的表示一张图片所有人脸框的二维数组过滤后转换为FaceInfo[]
     * @param dataOut 模型的输出，形状为[750, 6]的二维数组，其中每行表示 [class, probability, left, top, right, bottom]，其中最后4个为比例坐标
     * @param threshold 取值范围是0-1，表示只取置信度大于此数值的人脸
     * @returns 图片中所有人脸框识别的结果，未包含额外检测负载，返回的格式参考https://ai.baidu.com/sdk
     */
    postprocess(dataOut, threshold: number): FaceInfo[] {
        // data filter 筛掉人脸置信度小于阈值项
        const dataFilt = dataOut.filter(item => item[1] && item[1] > threshold);
        // data transform
        return dataFilt.map(item => this.transformData(item));
    }

    /**
     * 将模型的输出中的每一行，即长度为6的一维数组转换为FaceInfo
     * @param data 模型的输出，长度为6的一维数组，表示 [class, probability, left, top, right, bottom]，其中最后4个为比例坐标
     * @returns 一个人脸框识别的结果，未包含额外检测负载
     */
    transformData(data: Array<number>): FaceInfo {
        const face_info = {} as FaceInfo; // 存放输出结果
        const {fw, fh} = this.feedShape; // 给模型输入的宽高
        const {naturalWidth, naturalHeight} = this.inputSize; // 原始图片的宽高
        let dWidth: number; // 按shrink缩小后的宽度
        let dHeight: number; // 按shrink缩小后的高度

        if (naturalWidth > naturalHeight) {
            dHeight = this.scale;
            dWidth = dHeight * naturalWidth / naturalHeight;
        } else {
            dWidth = this.scale;
            dHeight = dWidth * naturalHeight / naturalWidth;
        }
        const dx = (fw - dWidth) / 2; // 左侧白边的宽度
        const dy = (fh - dHeight) / 2; // 上侧白边的高度

        face_info.face_probability = data[1]; // 人脸置信度
        // 在原图上的坐标
        face_info.location = {
            left: (data[2] * fw - dx) / dWidth * naturalWidth,
            top: (data[3] * fh - dy) / dHeight * naturalHeight,
            width: (data[4] - data[2]) * fw / dWidth * naturalWidth,
            height: (data[5] - data[3]) * fh / dHeight * naturalHeight
        } as Location;
        return face_info;
    }
}